{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "```shell\n",
    "pip install pymilvus==2.2.11\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import configparser\n",
    "import time\n",
    "import random\n",
    "from pymilvus import connections, utility\n",
    "from pymilvus import Collection, DataType, FieldSchema, CollectionSchema\n",
    "cfp = configparser.RawConfigParser()\n",
    "cfp.read('config_serverless.ini')\n",
    "milvus_uri = cfp.get('example', 'uri')\n",
    "token = cfp.get('example', 'token')\n",
    "\n",
    "connections.connect(\"default\",\n",
    "                    uri=milvus_uri,\n",
    "                    token=token)\n",
    "print(f\"Connecting to DB: {milvus_uri}\")\n",
    "\n",
    "# Check if the collection exists\n",
    "collection_name = \"book\"\n",
    "check_collection = utility.has_collection(collection_name)\n",
    "if check_collection:\n",
    "    drop_result = utility.drop_collection(collection_name)\n",
    "print(\"Success!\")\n",
    "# create a collection with customized primary field: book_id_field\n",
    "dim = 64\n",
    "book_id_field = FieldSchema(name=\"book_id\", dtype=DataType.INT64, is_primary=True, description=\"customized primary id\")\n",
    "word_count_field = FieldSchema(name=\"word_count\", dtype=DataType.INT64, description=\"word count\")\n",
    "book_intro_field = FieldSchema(name=\"book_intro\", dtype=DataType.FLOAT_VECTOR, dim=dim)\n",
    "schema = CollectionSchema(fields=[book_id_field, word_count_field, book_intro_field],\n",
    "                        auto_id=False,\n",
    "                        description=\"my first collection\")\n",
    "print(f\"Creating example collection: {collection_name}\")\n",
    "collection = Collection(name=collection_name, schema=schema)\n",
    "print(f\"Schema: {schema}\")\n",
    "print(\"Success!\")\n",
    "\n",
    "# insert data with customized ids\n",
    "nb = 1000\n",
    "insert_rounds = 2\n",
    "start = 0           # first primary key id\n",
    "total_rt = 0        # total response time for inert\n",
    "print(f\"Inserting {nb * insert_rounds} entities... \")\n",
    "for i in range(insert_rounds):\n",
    "    book_ids = [i for i in range(start, start+nb)]\n",
    "    word_counts = [random.randint(1, 100) for i in range(nb)]\n",
    "    book_intros = [[random.random() for _ in range(dim)] for _ in range(nb)]\n",
    "    entities = [book_ids, word_counts, book_intros]\n",
    "    t0 = time.time()\n",
    "    ins_resp = collection.insert(entities)\n",
    "    ins_rt = time.time() - t0\n",
    "    start += nb\n",
    "    total_rt += ins_rt\n",
    "print(f\"Succeed in {round(total_rt,4)} seconds!\")\n",
    "# print(f\"collection {collection_name} entities: {collection.num_entities}\")\n",
    "\n",
    "# flush\n",
    "print(\"Flushing...\")\n",
    "start_flush = time.time()\n",
    "collection.flush()\n",
    "end_flush = time.time()\n",
    "print(f\"Succeed in {round(end_flush - start_flush, 4)} seconds!\")\n",
    "# build index\n",
    "index_params = {\"index_type\": \"AUTOINDEX\", \"metric_type\": \"L2\", \"params\": {}}\n",
    "t0 = time.time()\n",
    "print(\"Building AutoIndex...\")\n",
    "collection.create_index(field_name=book_intro_field.name, index_params=index_params)\n",
    "t1 = time.time()\n",
    "print(f\"Succeed in {round(t1-t0, 4)} seconds!\")\n",
    "\n",
    "# load collection\n",
    "t0 = time.time()\n",
    "print(\"Loading collection...\")\n",
    "collection.load()\n",
    "t1 = time.time()\n",
    "print(f\"Succeed in {round(t1-t0, 4)} seconds!\")\n",
    "\n",
    "# search\n",
    "nq = 1\n",
    "search_params = {\"metric_type\": \"L2\",  \"params\": {\"level\": 2}}\n",
    "topk = 1\n",
    "for i in range(10):\n",
    "    search_vec = [[random.random() for _ in range(dim)] for _ in range(nq)]\n",
    "    print(f\"Searching vector: {search_vec}\")\n",
    "    t0 = time.time()\n",
    "    results = collection.search(search_vec,\n",
    "                            anns_field=book_intro_field.name,\n",
    "                            param=search_params,\n",
    "                            limit=topk,\n",
    "                            guarantee_timestamp=1)\n",
    "    t1 = time.time()\n",
    "    print(f\"Result:{results}\")\n",
    "    print(f\"search {i} latency: {round(t1-t0, 4)} seconds!\")\n",
    "\n",
    "connections.disconnect(\"default\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
